# -*- coding:utf-8 -*-

import json
import io
import os
import random
import tqdm
import re

def _make_w_io_base(f, mode: str):
    if not isinstance(f, io.IOBase):
        f_dirname = os.path.dirname(f)
        if f_dirname != "":
            os.makedirs(f_dirname, exist_ok=True)
        f = open(f, mode=mode)
    return f


def _make_r_io_base(f, mode: str):
    if not isinstance(f, io.IOBase):
        f = open(f, mode=mode)
    return f

def jdump(obj, f, mode="w", indent=4, default=str):
    """Dump a str or dictionary to a file in json format.

    Args:
        obj: An object to be written.
        f: A string path to the location on disk.
        mode: Mode for opening the file.
        indent: Indent for storing json dictionaries.
        default: A function to handle non-serializable entries; defaults to `str`.
    """
    f = _make_w_io_base(f, mode)
    if isinstance(obj, (dict, list)):
        json.dump(obj, f, indent=indent, default=default, ensure_ascii=False)
    elif isinstance(obj, str):
        f.write(obj)
    else:
        raise ValueError(f"Unexpected type: {type(obj)}")
    f.close()


def jload(f, mode="r"):
    """Load a .json file into a dictionary."""
    f = _make_r_io_base(f, mode)
    jdict = json.load(f)
    f.close()
    return jdict

def read_json(path):
    with open(path, "r", encoding="utf-8") as f:
        jdict = json.load(f)
    return jdict

def merge():

    d0 = read_json("./2199_0410.json")
    d1 = read_json("./2186_0410.json")
    new_data = []
    for instruct_dict in tqdm.tqdm(d0):
        instruct_dict_new = {}
        instruct_dict_new["instruction"] = instruct_dict.get("instruction", "")
        instruct_dict_new["input"] = instruct_dict.get("input", "")
        instruct_dict_new["output"] = instruct_dict.get("output", "")
        instruct_dict_new["instruction"] = "根据以下销售员与用户的对话，提取出用户在最后一次表述中的信息，包括意图、态度、实体、标签，每种信息以分号隔开，没有就输出为空，场景为黑牛保险。"
        new_data.append(instruct_dict_new)
    for instruct_dict in tqdm.tqdm(d1):
        instruct_dict_new = {}
        instruct_dict_new["instruction"] = instruct_dict.get("instruction", "")
        instruct_dict_new["input"] = instruct_dict.get("input", "")
        instruct_dict_new["output"] = instruct_dict.get("output", "")
        instruct_dict_new["instruction"] = "根据以下销售员与用户的对话，提取出用户在最后一次表述中的信息，包括意图、态度、实体、标签，每种信息以分号隔开，没有就输出为空，场景为360金融。"
        new_data.append(instruct_dict_new)

    random.shuffle(new_data)
    # print(new_data)
    jdump(new_data, "./2199-2186-0410.json")

def new_json():
    d = read_json("./train_0509_dialogue.json")
    # print(d)
    new_data = []
    for instruct_dict in tqdm.tqdm(d):
        instruct_dict_new = {}
        instruct_dict_new["instruction"] = instruct_dict.get("instruction", "")
        # if "保险" in instruct_dict_new["instruction"]:
        #     instruct_dict_new["instruction"] = "下面是段销售员和用户的对话记录。销售人员正在与用户进行赠险转短险的推广。需要根据对话中用户的最后一次表述进行信息的提取，信息包括三类，意图，态度，实体，任何一种信息没有就输出空。"
        # if "金融" in instruct_dict_new["instruction"]:
        #     instruct_dict_new["instruction"] = "下面是段销售员和用户的对话记录。销售人员正在向用户展示360金融产品的优势。需要根据对话中用户的最后一次表述进行信息的提取，信息包括三类，意图，态度，实体，任何一种信息没有就输出空。"

        if "保险" in instruct_dict_new["instruction"] or "赠险" in instruct_dict_new["instruction"]:
            instruct_dict_new["instruction"] = "下面是段销售员和用户的一次对话。销售人员在对用户进行赠险转短险的推广。需要对用户的话术进行信息提取，信息包括意图，态度，实体，任何一种信息没有就输出空。"
        if "金融" in instruct_dict_new["instruction"]:
            instruct_dict_new["instruction"] = "下面是段销售员和用户的一次对话。销售人员在向用户展示360金融产品的优势。需要对用户的话术进行信息提取，信息包括意图，态度，实体，任何一种信息没有就输出空。"

        instruct_dict_new["input"] = instruct_dict.get("input", "")
        idx = instruct_dict_new["input"].rfind("销售员:")
        if idx != -1:
            # idx = instruct_dict_new["input"].rfind("销售员:", 0, idx)
            # if idx != -1:
            instruct_dict_new["input"] = instruct_dict_new["input"][idx:]
            tmp = instruct_dict_new["input"].strip().split('\n')
            s = re.split(r'([，。！？])', tmp[0].strip())
            s = "销售员:" + "".join([i for i in s if i != ""][-2:])
            instruct_dict_new["input"] = s + "\n" + "\n".join(tmp[1:])
        instruct_dict_new["output"] = instruct_dict.get("output", "")
        new_data.append(instruct_dict_new)
    random.shuffle(new_data)
    jdump(new_data, "./train_0516_dialogue.json")

def length():
    d = read_json("./train_0517_dialogue.json")
    dct = {}
    for instruct_dict in tqdm.tqdm(d):
        string = str(instruct_dict)
        length = len(string)
        if length not in dct:
            dct[length] = []
        dct[length].append(string)
    lst = sorted(dct.items(), key=lambda x: x[0], reverse=True)
    for i in lst:
        if i[0] > 256:
            print(i[0], i[1])


if __name__ == "__main__":
    # merge()
    new_json()
    # length()